import tensorflow as tf
import os
import const

class MLPModel(object):
	def __init__(self, loadFromConst=False, modelDir=const.PATH_TO_DIR, diy = False) -> None:
		self.model = None
		self.history = None
		self.lastModelOutput = None

		if loadFromConst:
			load_file = modelDir if modelDir else const.PATH_TO_DIR
			self._loadExistModel(load_file)
		else:
			self._generateNewModel(diy=diy)
		
		self._printModelParams()

	def _printModelParams(self):
		if not self.model:
			exit(0)

		self.model.summary()
	
	def _generateNewModel(self, diy=const.USE_MLP_DIY):
		if self.model:
			return

		if not diy:
			self.model = tf.keras.Sequential([
				tf.keras.layers.Flatten(input_shape=const.INPUT_SHAPE),
				tf.keras.layers.Dense(const.LAYERS_1_NUM, activation=const.LAYERS_1_ACT),
				tf.keras.layers.Dense(const.LAYERS_2_NUM, activation=const.LAYERS_2_ACT), 
				tf.keras.layers.Dense(const.LAYERS_OUTPUT_NUM, activation=const.LAYERS_OUTPUT_ACT)
			])
		else:
			layerList = [None] * len(const.MLP_LAYER_DIY)
			for i, (attr, args) in enumerate(const.MLP_LAYER_DIY):
				f = getattr(tf.keras.layers, attr)
				try:
					if type(args) == type({}):
						layerList[i] = f(**args)
					if type(args) == type(()):
						layerList[i] = f(*args)
				except:
					print ('mlp自定义模型第%s层错误'%i)
					exit(0)
			self.model = tf.keras.Sequential(layerList)

		self.model.compile(**const.MODEL_COMPLIE_ARGS)

	def _loadExistModel(self, modelDir):
		if self.model:
			return
		load_file = modelDir+'model.h5'
		self.model = tf.keras.models.load_model(load_file)
		print ('MLP load from %s'%load_file)

	def trainModel(self, x, y, testRatio=0.05):
		if not self.model:
			return

		self.history = self.model.fit(x, y, batch_size=const.BATCH_SIZE, epochs=const.EPOCHS, validation_split=testRatio)
	
	def calcuModel(self, x):
		if not self.model:
			return
		
		self.lastModelOutput = self.model.predict(x)
		return self.lastModelOutput
	
	def saveModel(self, path):
		self.model.save(path + 'model.h5')
	
	def saveHistory(self, path):
		if not self.history:
			return

		history_path = path + 'history.json'
		history_dict = self.history.history
		with open(history_path, 'w') as file:
			file.write(str(history_dict))

	def getLastOneHistoryOnStr(self):
		return ('(%.3f-%.3f-%.3f-%.3f)'%(self.history.history['loss'][-1],self.history.history['accuracy'][-1],self.history.history['val_loss'][-1],self.history.history['val_accuracy'][-1]))

	def getLastOneHistoryOnFloat(self):
		return (self.history.history['loss'][-1],self.history.history['accuracy'][-1],self.history.history['val_loss'][-1],self.history.history['val_accuracy'][-1])

	def convertToTfLite(self, path=const.PATH_TO_DIR):
		if not self.model:
			return

		converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
		tflite_model = converter.convert()
		open(path+"model_tf.tflite", "wb").write(tflite_model)
		basic_model_size = os.path.getsize(path+"model_tf.tflite")
		print("Model is %d bytes" % basic_model_size)